# python .\photo\photo_num_in_zip.py --scan-dir D:\baseline\photo

from zipfile import ZipFile
from pathlib import Path
from argparse import ArgumentParser
from tqdm import tqdm

def analyze_zip_file(zip_path):
    file_stats = {}
    total_size = 0
    file_count = 0

    try:
        with ZipFile(zip_path, 'r') as archive:
            for file_info in archive.infolist():
                if not file_info.is_dir():
                    ext = Path(file_info.filename).suffix.lower()
                    if ext not in file_stats:
                        file_stats[ext] = {'count': 0, 'size': 0}
                    file_stats[ext]['count'] += 1
                    file_stats[ext]['size'] += file_info.file_size
                    total_size += file_info.file_size
                    file_count += 1
        return file_count, file_stats, total_size
    except Exception as e:
        print(f"分析ZIP文件 {zip_path} 时出错: {e}")
        return 0, {}, 0
    
def main():
    parser = ArgumentParser(description="分析指定目录下的所有ZIP文件")
    parser.add_argument('--scan-dir', required=True, help='要扫描的目录')
    
    args = parser.parse_args()
    scan_dir = Path(args.scan_dir)
    
    if not scan_dir.exists() or not scan_dir.is_dir():
        print(f"Error: 目录 {scan_dir} 不存在或不是一个目录")
        return
    
    zip_files = list(scan_dir.rglob('*.zip'))
    
    if not zip_files:
        print("未找到任何ZIP文件")
        return
    
    total_files_count = 0
    overall_file_stats = {}
    overall_size = 0
    
    with tqdm(total=len(zip_files), desc="处理ZIP文件") as pbar:
        for zip_file in zip_files:
            try:
                files_count, file_stats, size = analyze_zip_file(zip_file)
                
                total_files_count += files_count
                overall_size += size
                
                for ext, stats in file_stats.items():
                    if ext not in overall_file_stats:
                        overall_file_stats[ext] = {'count': 0, 'size': 0}
                    overall_file_stats[ext]['count'] += stats['count']
                    overall_file_stats[ext]['size'] += stats['size']
            except Exception as e:
                print(f"处理ZIP文件 {zip_file} 时出错: {e}")
            
            pbar.update(1)
    
    print("\n统计结果:")
    print(f"总共处理了 {len(zip_files)} 个ZIP文件.")
    print(f"总共包含 {total_files_count} 个文件.")
    print(f"总大小为 {overall_size / (1024**2):.2f} MB.")
    
    print("\n按文件类型统计:")
    for ext, stats in sorted(overall_file_stats.items()):
        print(f"{ext}: 文件数量 {stats['count']}, 总大小 {stats['size'] / (1024**2):.2f} MB")

if __name__ == "__main__":
    main()